"""
JWT认证中间件
提供JWT令牌验证、用户认证和权限管理功能
"""
import jwt
import time
import json
import hashlib
from typing import Optional, Dict, Any, List
from datetime import datetime, timedelta
from functools import wraps
from loguru import logger

from utils.constants import DEFAULT_DATETIME_FORMAT
from utils.logger import Logger


class JWTConfig:
    """JWT配置类"""
    
    def __init__(self):
        # 密钥配置
        self.secret_key = "your-secret-key-here"  # 应从环境变量读取
        self.algorithm = "HS256"
        self.access_token_expire_minutes = 30
        self.refresh_token_expire_days = 7
        
        # 安全配置
        self.token_blacklist_enabled = True
        self.max_concurrent_sessions = 5
        self.password_hash_iterations = 100000
        
        # Redis配置（用于令牌黑名单和会话管理）
        self.redis_key_prefix = "jwt:"
        self.blacklist_key = "blacklist"
        self.session_key = "session"


class JWTManager:
    """JWT管理器"""
    
    def __init__(self, redis_client, config: Optional[JWTConfig] = None):
        self.redis_client = redis_client
        self.config = config or JWTConfig()
        self.logger = Logger.get_logger(agent_name="jwt_manager")
        
    def create_access_token(self, user_id: str, username: str, 
                           permissions: List[str] = None, **kwargs) -> str:
        """创建访问令牌"""
        try:
            expire_time = datetime.utcnow() + timedelta(
                minutes=self.config.access_token_expire_minutes
            )
            
            payload = {
                "user_id": user_id,
                "username": username,
                "permissions": permissions or [],
                "token_type": "access",
                "exp": expire_time,
                "iat": datetime.utcnow(),
                "jti": self._generate_token_id(),
                **kwargs
            }
            
            token = jwt.encode(payload, self.config.secret_key, 
                             algorithm=self.config.algorithm)
            
            # 记录会话信息
            self._record_session(user_id, payload["jti"], expire_time)
            
            self.logger.info(f"访问令牌创建成功: user={username}, user_id={user_id}")
            return token
            
        except Exception as e:
            self.logger.error(f"访问令牌创建失败: {str(e)}")
            raise
    
    def create_refresh_token(self, user_id: str, username: str) -> str:
        """创建刷新令牌"""
        try:
            expire_time = datetime.utcnow() + timedelta(
                days=self.config.refresh_token_expire_days
            )
            
            payload = {
                "user_id": user_id,
                "username": username,
                "token_type": "refresh",
                "exp": expire_time,
                "iat": datetime.utcnow(),
                "jti": self._generate_token_id()
            }
            
            token = jwt.encode(payload, self.config.secret_key, 
                             algorithm=self.config.algorithm)
            
            self.logger.info(f"刷新令牌创建成功: user={username}")
            return token
            
        except Exception as e:
            self.logger.error(f"刷新令牌创建失败: {str(e)}")
            raise
    
    def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
        """验证令牌"""
        try:
            # 解码令牌
            payload = jwt.decode(token, self.config.secret_key, 
                               algorithms=[self.config.algorithm])
            
            # 检查令牌类型
            if payload.get("token_type") not in ["access", "refresh"]:
                self.logger.warning(f"无效令牌类型: {payload.get('token_type')}")
                return None
            
            # 检查黑名单
            if self._is_token_blacklisted(payload["jti"]):
                self.logger.warning(f"令牌在黑名单中: jti={payload['jti']}")
                return None
            
            # 检查会话有效性
            if not self._is_session_valid(payload["user_id"], payload["jti"]):
                self.logger.warning(f"会话无效: user_id={payload['user_id']}")
                return None
            
            self.logger.debug(f"令牌验证成功: user={payload['username']}")
            return payload
            
        except jwt.ExpiredSignatureError:
            self.logger.warning("令牌已过期")
            return None
        except jwt.InvalidTokenError as e:
            self.logger.warning(f"无效令牌: {str(e)}")
            return None
        except Exception as e:
            self.logger.error(f"令牌验证失败: {str(e)}")
            return None
    
    def refresh_access_token(self, refresh_token: str) -> Optional[str]:
        """刷新访问令牌"""
        try:
            payload = self.verify_token(refresh_token)
            if not payload or payload.get("token_type") != "refresh":
                return None
            
            # 创建新的访问令牌
            return self.create_access_token(
                user_id=payload["user_id"],
                username=payload["username"],
                permissions=payload.get("permissions", [])
            )
            
        except Exception as e:
            self.logger.error(f"刷新访问令牌失败: {str(e)}")
            return None
    
    def revoke_token(self, token: str) -> bool:
        """撤销令牌"""
        try:
            payload = self.verify_token(token)
            if not payload:
                return False
            
            # 添加到黑名单
            if self._add_to_blacklist(payload["jti"]):
                self.logger.info(f"令牌已撤销: jti={payload['jti']}, user={payload['username']}")
                return True
            
            return False
            
        except Exception as e:
            self.logger.error(f"令牌撤销失败: {str(e)}")
            return False
    
    def check_permission(self, token: str, required_permission: str) -> bool:
        """检查权限"""
        try:
            payload = self.verify_token(token)
            if not payload:
                return False
            
            user_permissions = payload.get("permissions", [])
            has_permission = required_permission in user_permissions
            
            self.logger.debug(f"权限检查: user={payload['username']}, "
                            f"permission={required_permission}, result={has_permission}")
            return has_permission
            
        except Exception as e:
            self.logger.error(f"权限检查失败: {str(e)}")
            return False
    
    def hash_password(self, password: str) -> str:
        """密码哈希"""
        try:
            # 使用PBKDF2进行密码哈希
            import hashlib
            salt = hashlib.sha256(str(time.time()).encode()).hexdigest()[:32]
            
            pwdhash = hashlib.pbkdf2_hmac(
                'sha256', 
                password.encode('utf-8'), 
                salt.encode('utf-8'), 
                self.config.password_hash_iterations
            ).hex()
            
            return f"{salt}${pwdhash}"
            
        except Exception as e:
            self.logger.error(f"密码哈希失败: {str(e)}")
            raise
    
    def verify_password(self, password: str, hashed: str) -> bool:
        """验证密码"""
        try:
            if "$" not in hashed:
                return False
            
            salt, pwdhash = hashed.split("$")
            
            pwdhash_new = hashlib.pbkdf2_hmac(
                'sha256',
                password.encode('utf-8'),
                salt.encode('utf-8'),
                self.config.password_hash_iterations
            ).hex()
            
            return pwdhash == pwdhash_new
            
        except Exception as e:
            self.logger.error(f"密码验证失败: {str(e)}")
            return False
    
    def _generate_token_id(self) -> str:
        """生成唯一令牌ID"""
        import uuid
        return str(uuid.uuid4())
    
    def _record_session(self, user_id: str, jti: str, expire_time: datetime):
        """记录会话信息"""
        try:
            session_key = f"{self.config.redis_key_prefix}{self.config.session_key}:{user_id}"
            
            # 获取当前会话列表
            sessions = self.redis_client.hgetall(session_key) or {}
            
            # 检查会话数量限制
            if len(sessions) >= self.config.max_concurrent_sessions:
                # 移除最旧的会话
                oldest_jti = min(sessions.keys(), key=lambda x: sessions[x])
                del sessions[oldest_jti]
            
            # 添加新会话
            sessions[jti] = expire_time.strftime(DEFAULT_DATETIME_FORMAT)
            
            # 保存到Redis
            self.redis_client.hsetall(session_key, sessions)
            
            # 设置过期时间
            ttl = int((expire_time - datetime.utcnow()).total_seconds())
            self.redis_client.expire(session_key, ttl)
            
        except Exception as e:
            self.logger.error(f"会话记录失败: {str(e)}")
    
    def _is_session_valid(self, user_id: str, jti: str) -> bool:
        """检查会话是否有效"""
        try:
            session_key = f"{self.config.redis_key_prefix}{self.config.session_key}:{user_id}"
            sessions = self.redis_client.hgetall(session_key) or {}
            
            return jti in sessions
            
        except Exception as e:
            self.logger.error(f"会话验证失败: {str(e)}")
            return False
    
    def _add_to_blacklist(self, jti: str) -> bool:
        """添加到黑名单"""
        try:
            blacklist_key = f"{self.config.redis_key_prefix}{self.config.blacklist_key}"
            
            # 设置黑名单过期时间为令牌过期时间 + 1小时
            expire_time = datetime.utcnow() + timedelta(hours=1)
            ttl = int((expire_time - datetime.utcnow()).total_seconds())
            
            self.redis_client.hset(blacklist_key, jti, "1")
            self.redis_client.expire(blacklist_key, ttl)
            
            return True
            
        except Exception as e:
            self.logger.error(f"添加到黑名单失败: {str(e)}")
            return False
    
    def _is_token_blacklisted(self, jti: str) -> bool:
        """检查令牌是否在黑名单中"""
        try:
            blacklist_key = f"{self.config.redis_key_prefix}{self.config.blacklist_key}"
            return self.redis_client.hexists(blacklist_key, jti)
            
        except Exception as e:
            self.logger.error(f"黑名单检查失败: {str(e)}")
            return False


def jwt_required(permissions: List[str] = None):
    """JWT认证装饰器"""
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 这里需要从请求中获取token，具体实现取决于使用的Web框架
            # 示例：token = request.headers.get("Authorization", "").replace("Bearer ", "")
            
            # 临时实现：从kwargs中获取token（用于测试）
            token = kwargs.get("token", "")
            if not token:
                raise ValueError("缺少认证令牌")
            
            # 这里需要注入JWT管理器实例
            jwt_manager = kwargs.get("jwt_manager")
            if not jwt_manager:
                raise ValueError("缺少JWT管理器")
            
            # 验证令牌
            payload = jwt_manager.verify_token(token)
            if not payload:
                raise ValueError("无效或过期的令牌")
            
            # 检查权限
            if permissions:
                user_permissions = payload.get("permissions", [])
                if not any(perm in user_permissions for perm in permissions):
                    raise ValueError("权限不足")
            
            # 添加用户信息到kwargs
            kwargs["current_user"] = payload
            
            return func(*args, **kwargs)
        return wrapper
    return decorator